import torch
from tqdm import tqdm
from torch.utils.data import RandomSampler,DataLoader
import torch.nn as nn
from transformers import BertModel, BertTokenizer,BertForMaskedLM,BertConfig
import copy,re,pickle,random
import numpy as np

with open(r"med_map.pkl",'rb') as file:
    med_map=pickle.load(file)
key = ['食盐','石盐','盐绿','炒盐','井盐','大青盐','烧盐','沧盐']

# for item in med_map:
#     if item in key: 
#         med_map.update({item:'盐粒'})
# with open(r"med_map.pkl",'wb') as file:
#     pickle.dump(med_map, file)

# 把盐改为盐粒
    
for item in med_map:
    if item in key: 
        print(med_map[item])






with open("standard_herb_list.txt",'r',encoding='utf-8') as file:
    content = file.readlines()
    herbs = [i.strip() for i in content]
print("len(herbs):", len(herbs))

count = 0
with open('val.txt', 'r', encoding='utf-8') as file:
    content = file.readlines()
    for item in content:
        sent = item.strip().split('\t\t')
        # print(sent)
        for item in re.split('[\[\]]', sent[1]):
            if item.strip('|') in med_map and med_map[item.strip('|')] in herbs:  # 如果修正之后是标准药名
                count += 1

print(count)
for k,v in med_map.items():
    if v not in herbs + ['盐']:
        print(k,v)
    if v == '茯苓':
        print(k)
        
        
tokenizer = BertTokenizer.from_pretrained('../chinese_L-12_H-768_A-12')
print(tokenizer.convert_tokens_to_ids('盐')) 
print(tokenizer.convert_tokens_to_ids('乌头'))
herb_list = list(range(21128, 21662)) + [4663]

config = BertConfig.from_pretrained('../chinese_L-12_H-768_A-12')

print('1',config.vocab_size)
print(len(tokenizer.vocab))
random.seed(2021)

import random
print(random.choice(herbs) )
print(random.choice(herbs) )
print(random.choice(herbs) )


